// CreateInitrd.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <windows.h>
#include "romdisk.h"

FILE *romdisk_file;

typedef struct RD_RELOCATION {
	uint32_t relocation_pos;
	uint32_t relocation_value;
} RD_RELOCATION_T;

RD_RELOCATION_T relocation_table[4];//0 - image size,1 - list ptr,2 - data ptr,3 - fcount

bool log_progress = true;

void relocate(RD_RELOCATION_T *reloc){
	size_t pos = ftell(romdisk_file);
	fseek(romdisk_file,reloc->relocation_pos,SEEK_SET);
	fwrite(&(reloc->relocation_value),sizeof(uint32_t),1,romdisk_file);
	fseek(romdisk_file,pos,SEEK_SET);
}

bool add_file(char *file,char *name){
	FILE *f;
	size_t length;
	ROMDISK_FILE_T fileinfo;
	char *buf;
	if (log_progress)
		printf("Adding %s\n",file);
	f = fopen(file,"rb");
	if (f == NULL){
		printf("Error: Could not open input file %s\n",file);
		return false;
	}		
	//Seek to end of file
	fseek(f,0,SEEK_END);
	length = ftell(f);//Poll file position
	rewind(f);//Seek to 0
	//Write start address
	relocation_table[1].relocation_value = (uint32_t) ftell(romdisk_file);
	relocate(&relocation_table[1]);
	//First dword is next ptr so this works
	relocation_table[1].relocation_pos = relocation_table[1].relocation_value;
	relocation_table[2].relocation_pos = relocation_table[1].relocation_value + 4;//Data ptr
	//Fill header
	strcpy((char *)fileinfo.filename,name);
	fileinfo.length = (uint32_t) length;
	//Write header
	if (sizeof (ROMDISK_FILE_T) != fwrite(&fileinfo,1,sizeof(ROMDISK_FILE_T),romdisk_file)){
		printf("Error: Could write full header to romdisk %s\n",file);
		return false;
	}		
	//Write data address
	relocation_table[2].relocation_value = (uint32_t) ftell(romdisk_file);
	relocate(&relocation_table[2]);
	//Write actual data
	buf = (char *) malloc(length);
	length = fread(buf,1,length,f);
	if (length != fwrite(buf,1,length,romdisk_file)){
		printf("Error: Could write full file to romdisk %s\n",file);
		return false;
	}		
	//Clean up the mess we made
	free(buf);
	fclose(f);
	relocation_table[3].relocation_value++;
	return true;
}


bool iterate_files(char *dir){
	HANDLE hFind = INVALID_HANDLE_VALUE;
	WIN32_FIND_DATA ffd;
	char *spec = (char *)malloc(strlen(dir)+4);
	strcpy(spec,dir);
	strcat(spec,"\\*");
	hFind = FindFirstFile(spec, &ffd);
	if (hFind == INVALID_HANDLE_VALUE)  {
		printf("Error: Directory %s does not exist\n",dir);
		free((void*)spec);
		return false;
	} 
	do {
		if (strcmp(ffd.cFileName, ".") != 0 && 
			strcmp(ffd.cFileName, "..") != 0) {
				if (!(ffd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)){
					char *path = (char *) malloc(strlen(ffd.cFileName) + strlen(dir) + 3);
					strcpy(path,dir);
					strcat(path,"\\");
					strcat(path,ffd.cFileName);
					if (!add_file(path,ffd.cFileName)){
						FindClose(hFind);
						return false;
					} else {
						free((void*) path);
					}
				}
		}
	} while (FindNextFile(hFind, &ffd) != 0);
	if (GetLastError() != ERROR_NO_MORE_FILES) {
		FindClose(hFind);
		printf("Error iterating directory %s\n",dir);
		free((void*)spec);
		return false;
	}
	FindClose(hFind);
	hFind = INVALID_HANDLE_VALUE;
	free((void*)spec);
	return true;
}


int _tmain(int argc, char* argv[])
{
	ROMDISK_HEADER_T rd_header;
	printf("CreateInitrd by Peterbjornx %s %i\n",argv[2],argc);
	if (argc < 2){
		printf("Usage: createinitrd <filename> <directory> [label]\n");
		return 1;
	}
	romdisk_file = fopen(argv[1],"wb");
	if (romdisk_file == NULL){
		printf("Error: Can not create image %s\n",argv[1]);
		return 1;
	}
	rd_header.magic = ROMDISK_MAGIC;
	if (argc >= 3)
		strcpy((char *)rd_header.volname,argv[3]);
	else
		strcpy((char *)rd_header.volname,"romdisk");
	relocation_table[1].relocation_pos = 4;
	relocation_table[0].relocation_pos = 8;
	relocation_table[3].relocation_pos = 12;
	relocation_table[3].relocation_value = 0;
	long written = fwrite(&rd_header,1,sizeof(ROMDISK_HEADER_T),romdisk_file);
	if (written != sizeof(ROMDISK_HEADER_T)){
		printf("Error: Can not write image header %i %i\n",written,sizeof(ROMDISK_HEADER_T));
		fclose(romdisk_file);
		return 1;
	}
	if (!iterate_files(argv[2])){
		fclose(romdisk_file);
		return 1;
	}
	relocation_table[0].relocation_value = (uint32_t) ftell(romdisk_file);
	relocation_table[1].relocation_value = 0x12345;
	relocate(&relocation_table[0]);//Size
	relocate(&relocation_table[1]);//End of list
	relocate(&relocation_table[3]);//File count
	fclose(romdisk_file);	
	return 0;
}

